import math
import random
import numpy as np

from PIL import Image
from PIL import ImageFilter

# import torchvision.transforms.functional as TF
# from torchvision import transforms
# import torch
import mindspore as ms
from mindspore import Tensor
from mindspore.dataset import vision
from mindspore import ops
import numpy as np
from PIL import Image, ImageEnhance
# import torchvision.transforms.functional as TF
# from torchvision import transforms
import numbers
from dataset.color_jitter import ColorJitter


def _to_norm(imgs, labels):
    # to tensor

    imgs = [np.array(img, np.float32).transpose(2, 0, 1) for img in imgs]
    labels_tensor = [np.array(label, np.int32) for label in labels]
    # print(labels_tensor)

    mean = 127.5
    std = 127.5

    normalize_imgs = [(img - mean)/std for img in imgs]


    return normalize_imgs, labels_tensor


class CDDataAugmentation:
    def __init__(
            self,
            img_size,
            with_random_hflip=False,
            with_random_vflip=False,
            with_random_rot=False,
            with_random_crop=False,
            with_scale_random_crop=False,
            with_random_blur=False,
            random_color_tf=False
    ):
        self.img_size = img_size
        if self.img_size is None:
            self.img_size_dynamic = True
        else:
            self.img_size_dynamic = False
        self.with_random_hflip = with_random_hflip
        self.with_random_vflip = with_random_vflip
        self.with_random_rot = with_random_rot
        self.with_random_crop = with_random_crop
        self.with_scale_random_crop = with_scale_random_crop
        self.with_random_blur = with_random_blur
        self.random_color_tf=random_color_tf
    def transform(self, imgs, labels, to_normalize=True):
        """
        :param imgs: [ndarray,]
        :param labels: [ndarray,]
        :return: [ndarray,],[ndarray,]
        """
        # resize image and covert to tensor


        if not self.img_size_dynamic:
            if imgs[0].size != (self.img_size, self.img_size):
                imgs = [img.resize((self.img_size, self.img_size), Image.BICUBIC) for img in imgs]
        else:
            self.img_size = imgs[0].size[0]

        labels = [Image.fromarray(img) for img in labels]
        if len(labels) != 0:
            if labels[0].size != (self.img_size, self.img_size):
                labels = [label.resize((self.img_size, self.img_size), Image.NEAREST) for label in labels]

        random_base = 0.5
        if self.with_random_hflip and random.random() > 0.5:
            imgs = [hflip_image(img) for img in imgs]
            labels = [hflip_image(img) for img in labels]

        if self.with_random_vflip and random.random() > 0.5:
            imgs = [vflip_image(img) for img in imgs]
            labels = [vflip_image(img) for img in labels]

        if self.with_random_rot and random.random() > random_base:
            angles = [90, 180, 270]
            index = random.randint(0, 2)
            angle = angles[index]
            imgs = [rotate_img(img, angle) for img in imgs]
            labels = [rotate_img(img, angle) for img in labels]

        if self.with_random_crop and random.random() > 0:
            imgs = [random_crop_resize_img(img=img, target_size=(self.img_size, self.img_size),scale=[0.8, 1.2],
                                           ratio=[1., 1.], interpolation=Image.CUBIC) for img in imgs]

            labels = [random_crop_resize_img(label, (self.img_size, self.img_size),(0.8, 1.2),
                                           interpolation=Image.CUBIC) for label in labels]

        if self.with_scale_random_crop:
            # rescale
            scale_range = [1, 1.2]
            target_scale = scale_range[0] + random.random() * (scale_range[1] - scale_range[0])

            imgs = [pil_rescale(img, target_scale, order=3) for img in imgs]
            labels = [pil_rescale(img, target_scale, order=0) for img in labels]
            # crop
            imgsize = imgs[0].size  # h, w
            box = get_random_crop_box(imgsize=imgsize, cropsize=self.img_size)
            imgs = [pil_crop(img, box, cropsize=self.img_size, default_value=0)
                    for img in imgs]
            labels = [pil_crop(img, box, cropsize=self.img_size, default_value=255)
                    for img in labels]

        if self.with_random_blur and random.random() > 0:
            radius = random.random()
            imgs = [img.filter(ImageFilter.GaussianBlur(radius=radius))
                    for img in imgs]

        if self.random_color_tf:
            color_jitter = ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3)
            imgs = [color_jitter.forward(img) for img in imgs]
            
        if to_normalize:
            imgs_tensor, labels_tensor = _to_norm(imgs, labels)

        return imgs_tensor, labels_tensor


def pil_crop(image, box, cropsize, default_value):
    assert isinstance(image, Image.Image)
    img = np.array(image)

    if len(img.shape) == 3:
        cont = np.ones((cropsize, cropsize, img.shape[2]), img.dtype)*default_value
    else:
        cont = np.ones((cropsize, cropsize), img.dtype)*default_value
    cont[box[0]:box[1], box[2]:box[3]] = img[box[4]:box[5], box[6]:box[7]]

    return Image.fromarray(cont)


def get_random_crop_box(imgsize, cropsize):
    h, w = imgsize
    ch = min(cropsize, h)
    cw = min(cropsize, w)

    w_space = w - cropsize
    h_space = h - cropsize

    if w_space > 0:
        cont_left = 0
        img_left = random.randrange(w_space + 1)
    else:
        cont_left = random.randrange(-w_space + 1)
        img_left = 0

    if h_space > 0:
        cont_top = 0
        img_top = random.randrange(h_space + 1)
    else:
        cont_top = random.randrange(-h_space + 1)
        img_top = 0

    return cont_top, cont_top+ch, cont_left, cont_left+cw, img_top, img_top+ch, img_left, img_left+cw


def pil_rescale(img, scale, order):
    assert isinstance(img, Image.Image)
    height, width = img.size
    target_size = (int(np.round(height*scale)), int(np.round(width*scale)))
    return pil_resize(img, target_size, order)


def pil_resize(img, size, order):
    assert isinstance(img, Image.Image)
    if size[0] == img.size[0] and size[1] == img.size[1]:
        return img
    if order == 3:
        resample = Image.BICUBIC
    elif order == 0:
        resample = Image.NEAREST
    return img.resize(size[::-1], resample)

def hflip_image(img):
    img = img.transpose(Image.FLIP_LEFT_RIGHT)
    return img

def vflip_image(img):
    img = img.transpose(Image.FLIP_TOP_BOTTOM)
    return img

def rotate_img(img, angle):
    img = img.rotate(angle)
    return img

def random_crop_resize_img(img, target_out_size, scale=[0.08, 1.0], ratio=[3. / 4., 4. / 3.], mode = Image.BILINEAR):
    aspect_ratio = math.sqrt(np.random.uniform(*ratio))
    w = 1. * aspect_ratio
    h = 1. / aspect_ratio

    bound = min((float(img.size[0]) / img.size[1]) / (w**2),
                (float(img.size[1]) / img.size[0]) / (h**2))
    scale_max = min(scale[1], bound)
    scale_min = min(scale[0], bound)

    target_area = img.size[0] * img.size[1] * np.random.uniform(scale_min,
                                                                scale_max)
    target_size = math.sqrt(target_area)
    w = int(target_size * w)
    h = int(target_size * h)

    i = np.random.randint(0, img.size[0] - w + 1)
    j = np.random.randint(0, img.size[1] - h + 1)

    img = img.crop((i, j, i + w, j + h))
    img = img.resize((int(target_out_size), int(target_out_size)), mode)
    return img




